import json
import os
import random
import re
import ast

def jsonl_to_random_json(jsonl_file_path, output_dir, num_files):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 读取.jsonl文件中的所有行
    with open(jsonl_file_path, 'r') as jsonl_file:
        lines = [json.loads(line) for line in jsonl_file]

    # 随机选择num_files数量的行
    random_lines = random.sample(lines, min(num_files, len(lines)))
    
    # 将随机选择的JSON对象写入单独的.json文件中
    for i, json_obj in enumerate(random_lines):
        # 构建输出文件的路径
        file_path = os.path.join(output_dir, f'output_{i+1}.json')
        
        # 将JSON对象写入到单独的.json文件中
        with open(file_path, 'w') as json_file:
            json.dump(json_obj, json_file)
            json_file.write('\n')  # 确保每个JSON对象占一行


def jsonl_to_random_json_and_split(jsonl_file_path, output_dir, num_files):
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 读取.jsonl文件中的所有行
    with open(jsonl_file_path, 'r', encoding='utf-8') as jsonl_file:
        lines = [json.loads(line) for line in jsonl_file]

    # 过滤掉difficulty属性值为introductory的JSON对象
    filtered_lines = [line for line in lines if line.get('difficulty') == 'competition']

    # 随机选择num_files数量的行
    random_lines = random.sample(filtered_lines, min(num_files, len(filtered_lines)))
    
    # 将随机选择的JSON对象写入单独的.json文件中
    for i, json_obj in enumerate(random_lines):
        # 构建输出文件的路径，文件名直接使用数字
        json_file_path = os.path.join(output_dir, f'{i+1}.json')

        # 将JSON对象写入到单独的.json文件中
        with open(json_file_path, 'w', encoding='utf-8') as json_file:
            json.dump(json_obj, json_file)
            json_file.write('\n')  # 确保每个JSON对象占一行
        
        # 创建以.json文件名命名的文件夹
        attribute_dir = os.path.join(output_dir, f'{i+1}')
        if not os.path.exists(attribute_dir):
            os.makedirs(attribute_dir)
        
        # 解析JSON对象并拆分属性到单独的文件
        for attribute in ['question', 'input_output']:
            # 构建属性文件的路径，文件名以属性值命名
            attribute_value = json_obj.get(attribute, '')
            attribute_file_path = os.path.join(attribute_dir, f"{attribute}_{i+1}.txt")
            
            # 将属性值写入到单独的文件中
            with open(attribute_file_path, 'w', encoding='utf-8') as attribute_file:
                attribute_file.write(str(attribute_value))

        # 特殊处理'solutions'属性
        solutions = json_obj.get('solutions')
        t_s = ast.literal_eval(solutions)
        solutions_dir = os.path.join(attribute_dir, 'solutions')
        if not os.path.exists(solutions_dir):
            os.makedirs(solutions_dir)
        
        # 将'solutions'数组中的每个字符串保存为一个新的文件
        j = 0
        for solution in t_s:
            # 处理转义符
            # processed_solution = solution.encode().decode('unicode_escape')
            solution_file_path = os.path.join(solutions_dir, f"solution_{j+1}.py")

            
            j = j + 1
            
            # 将处理后的代码写入到单独的文件中
            with open(solution_file_path, 'w', encoding='utf-8') as solution_file:
                solution_file.write(solution)



# 使用示例
jsonl_file_path = 'C:\\data\\study\\postgraduate\\ldm\\dataset\\datasets--codeparrot--apps\\snapshots\\21e74ddf8de1a21436da12e3e653065c5213e9d1\\train.jsonl'  # 你的.jsonl文件路径
output_dir = 'C:\\data\\study\\postgraduate\\ldm\\dataset\\datasets--codeparrot--apps\\data'           # 输出目录路径
max_json_files = 10                            # 最大生成的.json文件数量

jsonl_to_random_json_and_split(jsonl_file_path, output_dir, max_json_files)